{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2025-08-03T08:25:49.299814Z",
     "start_time": "2025-08-03T08:25:49.292563Z"
    }
   },
   "source": "from init_llm_model import llm",
   "outputs": [],
   "execution_count": 3
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-03T08:25:49.327351Z",
     "start_time": "2025-08-03T08:25:49.315448Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_openai import ChatOpenAI\n",
    "from pydantic import BaseModel, Field\n",
    "\n",
    "tagging_prompt = ChatPromptTemplate.from_template(\n",
    "    \"\"\"\n",
    "Extract the desired information from the following passage.\n",
    "\n",
    "Only extract the properties mentioned in the 'Classification' function.\n",
    "\n",
    "Passage:\n",
    "{input}\n",
    "\"\"\"\n",
    ")\n",
    "\n",
    "\n",
    "class Classification(BaseModel):\n",
    "    sentiment: str = Field(description=\"The sentiment of the text\")\n",
    "    aggressiveness: int = Field(\n",
    "        description=\"How aggressive the text is on a scale from 1 to 10\"\n",
    "    )\n",
    "    language: str = Field(description=\"The language the text is written in\")\n",
    "\n",
    "\n",
    "# Structured LLM\n",
    "structured_llm = llm.with_structured_output(Classification)"
   ],
   "id": "881ad46a7a2054f4",
   "outputs": [],
   "execution_count": 4
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-03T08:25:54.035309Z",
     "start_time": "2025-08-03T08:25:49.341884Z"
    }
   },
   "cell_type": "code",
   "source": [
    "inp = \"Estoy increiblemente contento de haberte conocido! Creo que seremos muy buenos amigos!\"\n",
    "prompt = tagging_prompt.invoke({\"input\": inp})\n",
    "response = structured_llm.invoke(prompt)\n",
    "\n",
    "print(response)"
   ],
   "id": "98662d736ef111d",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sentiment='positive' aggressiveness=1 language='Spanish'\n"
     ]
    }
   ],
   "execution_count": 5
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "If we want dictionary output, we can just call .model_dump()",
   "id": "dad97d68dc6b74a4"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-03T08:25:58.335402Z",
     "start_time": "2025-08-03T08:25:54.169723Z"
    }
   },
   "cell_type": "code",
   "source": [
    "inp = \"Estoy muy enojado con vos! Te voy a dar tu merecido!\"\n",
    "prompt = tagging_prompt.invoke({\"input\": inp})\n",
    "response = structured_llm.invoke(prompt)\n",
    "\n",
    "response.model_dump()"
   ],
   "id": "29befe5175d7b388",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'sentiment': 'enojado', 'aggressiveness': 8, 'language': 'es'}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 6
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "Careful schema definition gives us more control over the model's output.\n",
    "\n",
    "Specifically, we can define:\n",
    "\n",
    "Possible values for each property\n",
    "Description to make sure that the model understands the property\n",
    "Required properties to be returned"
   ],
   "id": "504803c265bea868"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-03T08:25:58.461273Z",
     "start_time": "2025-08-03T08:25:58.444506Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class Classification(BaseModel):\n",
    "    sentiment: str = Field(..., enum=[\"happy\", \"neutral\", \"sad\"])\n",
    "    aggressiveness: int = Field(..., description=\"describes how aggressive the statement is, the higher the number the more aggressive\",\n",
    "        enum=[1, 2, 3, 4, 5],\n",
    "    )\n",
    "    language: str = Field(\n",
    "        ..., enum=[\"spanish\", \"english\", \"french\", \"german\", \"italian\"]\n",
    "    )"
   ],
   "id": "3949bca61eb39b05",
   "outputs": [],
   "execution_count": 7
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-03T08:25:58.490684Z",
     "start_time": "2025-08-03T08:25:58.481803Z"
    }
   },
   "cell_type": "code",
   "source": [
    "tagging_prompt = ChatPromptTemplate.from_template(\n",
    "    \"\"\"\n",
    "Extract the desired information from the following passage.\n",
    "\n",
    "Only extract the properties mentioned in the 'Classification' function.\n",
    "\n",
    "Passage:\n",
    "{input}\n",
    "\"\"\"\n",
    ")\n",
    "\n",
    "llm = llm.with_structured_output(Classification)"
   ],
   "id": "1de59ab33a14341d",
   "outputs": [],
   "execution_count": 8
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-03T08:26:30.673342Z",
     "start_time": "2025-08-03T08:26:26.142984Z"
    }
   },
   "cell_type": "code",
   "source": [
    "inp = \"Estoy increiblemente contento de haberte conocido! Creo que seremos muy buenos amigos!\"\n",
    "prompt = tagging_prompt.invoke({\"input\": inp})\n",
    "llm.invoke(prompt)"
   ],
   "id": "1dca77e948c0838b",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Classification(sentiment='happy', aggressiveness=1, language='spanish')"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 9
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
